import datetime
import sys
import threading
import tkinter as tk
from tkinter import ttk, messagebox, filedialog,simpledialog
from ttkthemes import ThemedTk
import random
import string
import uuid
import pymysql
import cx_Oracle
import psycopg2
from db_generdata_bigdata import create_table
import hashlib
from Table_tkinter import TextRedirector
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
import base64


# oracle_path = os.path.join(os.path.dirname(__file__), "instantclient_21_17")
# cx_Oracle.init_oracle_client(lib_dir='./table_generator_env/Lib/site-packages/instantclient_21_17;table_generator_env/Lib/site-packages/instantclient_21_17')

class TableGeneratorApp:
    def __init__(self, root):
        # self.root = ThemedTk(theme="arc")
        self.root = root
        self.root.title("数据库建表与数据生成工具")
        self.root.geometry("1000x800")
        self.clicked_build_encrypted_data = False  # 新增标志变量
        self.setup_ui()

    def setup_ui(self):
        style = ttk.Style()
        style.configure("TLabel", font=("Arial", 11))
        style.configure("TEntry", font=("Arial", 11))
        style.configure("TButton", font=("Arial", 11))
        style.configure("TCheckbutton", font=("Arial", 11))

        db_frame = ttk.LabelFrame(self.root, text="数据库连接", padding=10)
        db_frame.pack(fill=tk.X, padx=10, pady=5)

        # 数据库连接参数部分
        db_frame = ttk.LabelFrame(root, text="数据库连接", padding=10)
        db_frame.pack(fill=tk.X, padx=10, pady=5)

        # 第一行参数
        ttk.Label(db_frame, text="数据库类型:").grid(row=0, column=0, padx=5, pady=5, sticky=tk.E)
        self.db_type = ttk.Combobox(db_frame,
                                    values=["Oracle", "MySQL", "PostgreSQL", "GoldenDB", "OceanBase", "DM", "TDSQL"],
                                    state="readonly", width=10)
        self.db_type.grid(row=0, column=1, padx=5, pady=5, sticky=tk.W)
        self.db_type.current(0)

        ttk.Label(db_frame, text="主机:").grid(row=0, column=2, padx=5, pady=5, sticky=tk.E)
        self.host_entry = ttk.Entry(db_frame, width=18)
        self.host_entry.insert(0, "10.6.68.183")
        self.host_entry.grid(row=0, column=3, padx=5, pady=5, sticky=tk.W)

        ttk.Label(db_frame, text="端口:").grid(row=0, column=4, padx=5, pady=5, sticky=tk.E)
        self.port_entry = ttk.Entry(db_frame, width=8)
        self.port_entry.insert(0, "1521")
        self.port_entry.grid(row=0, column=5, padx=5, pady=5, sticky=tk.W)

        # 第二行参数
        ttk.Label(db_frame, text="数据库/实例名:").grid(row=1, column=0, padx=5, pady=5, sticky=tk.E)
        self.dbname_entry = ttk.Entry(db_frame, width=10)
        self.dbname_entry.insert(0, "leora")
        self.dbname_entry.grid(row=1, column=1, padx=5, pady=5, sticky=tk.W)

        ttk.Label(db_frame, text="Schema:").grid(row=1, column=2, padx=5, pady=5, sticky=tk.E)
        self.schema_entry = ttk.Entry(db_frame, width=10)
        self.schema_entry.insert(0, "ADMS")
        self.schema_entry.grid(row=1, column=3, padx=5, pady=5, sticky=tk.W)

        ttk.Label(db_frame, text="用户名:").grid(row=1, column=4, padx=5, pady=5, sticky=tk.E)
        self.user_entry = ttk.Entry(db_frame, width=10)
        self.user_entry.insert(0, "system")
        self.user_entry.grid(row=1, column=5, padx=5, pady=5, sticky=tk.W)

        ttk.Label(db_frame, text="密码:").grid(row=1, column=6, padx=5, pady=5, sticky=tk.E)
        self.pwd_entry = ttk.Entry(db_frame, show="*", width=10)
        self.pwd_entry.insert(0, "root")
        self.pwd_entry.grid(row=1, column=7, padx=5, pady=5, sticky=tk.W)

        # 参数配置区域
        frame = ttk.LabelFrame(self.root, text="表生成参数", padding=10)
        frame.pack(fill=tk.X, padx=10, pady=5)

        self.table_count_entry = ttk.Entry(frame, width=10)
        self.table_count_entry.insert(0, "10")
        self.table_count_entry.grid(row=0, column=3)
        ttk.Label(frame, text="表数量:").grid(row=0, column=2, sticky=tk.W)

        self.columns_entry = ttk.Entry(frame, width=10)
        self.columns_entry.insert(0, "10")
        self.columns_entry.grid(row=1, column=1)
        ttk.Label(frame, text="字段数:").grid(row=1, column=0, sticky=tk.W)

        self.rows_entry = ttk.Entry(frame, width=10)
        self.rows_entry.insert(0, "2000")
        self.rows_entry.grid(row=1, column=3)
        ttk.Label(frame, text="插入行数:").grid(row=1, column=2, sticky=tk.W)

        self.col_prefix_entry = ttk.Entry(frame, width=10)
        self.col_prefix_entry.insert(0, "col")
        self.col_prefix_entry.grid(row=2, column=1)
        ttk.Label(frame, text="字段前缀:").grid(row=2, column=0, sticky=tk.W)

        self.table_prefix_entry = ttk.Entry(frame, width=10)
        self.table_prefix_entry.insert(0, "b")
        self.table_prefix_entry.grid(row=0, column=1)
        ttk.Label(frame, text="表前缀:").grid(row=0, column=0, sticky=tk.W)

        self.col_len = ttk.Entry(frame, width=10)
        self.col_len.insert(0, "50")
        self.col_len.grid(row=2, column=3)
        ttk.Label(frame, text="字符长度:").grid(row=2, column=2, sticky=tk.W)
        

        self.auto_number_var = tk.BooleanVar(value=True)
        self.add_pk_var = tk.BooleanVar(value=True)
        self.random_name_var = tk.BooleanVar(value=False)
        self.show_type_comment_var = tk.BooleanVar(value=False)
        self.create_large_table_var = tk.BooleanVar(value=False)
        self.show_type_comment_var = tk.BooleanVar(value=False)
        self.build_encrypt_data_var = tk.BooleanVar(value=False)
        self.encrypt_md5_var = tk.BooleanVar(value=False)
        self.encrypt_aes_var = tk.BooleanVar(value=False)
        self.encrypt_sha256_var = tk.BooleanVar(value=False)
        self.encrypt_type_checkbuttons = []
        # 现有的 Checkbuttons
        ttk.Checkbutton(frame, text="表名自动编号", variable=self.auto_number_var).grid(row=3, column=0)
        ttk.Checkbutton(frame, text="添加主键", variable=self.add_pk_var).grid(row=3, column=1)
        ttk.Checkbutton(frame, text="表名随机", variable=self.random_name_var).grid(row=3, column=2)
        ttk.Checkbutton(frame, text="字段类型注释", variable=self.show_type_comment_var).grid(row=3, column=3)
        build_encrypt_cb = ttk.Checkbutton(frame, text="构建加密数据", variable=self.build_encrypt_data_var, command=lambda: self.toggle_encrypt_type(frame))
        build_encrypt_cb.grid(row=3, column=6)
        # 按钮区域
        button_frame = ttk.Frame(self.root, padding=10)
        button_frame.pack(fill=tk.X)

        ttk.Button(button_frame, text="生成预览", command=self.generate_sql).pack(side=tk.LEFT)
        ttk.Button(button_frame, text="执行SQL (模拟)", command=self.simulate_execute).pack(side=tk.LEFT, padx=5)
        ttk.Button(button_frame, text="导出SQL", command=self.export_sql).pack(side=tk.LEFT, padx=5)
        ttk.Button(button_frame, text="清空", command=self.clear_output).pack(side=tk.LEFT)
        ttk.Button(button_frame, text="测试连接", command=self.test_connection).pack(side=tk.LEFT, padx=5)
        # ttk.Button(button_frame, text="构建加密数据", command=self.build_encrypted_data).pack(side=tk.LEFT, padx=5)

        self.output_text = tk.Text(self.root, height=30, wrap=tk.NONE, font=("Courier New", 10))
        self.output_text.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        self.output_redirector = TextRedirector(self.output_text)
        sys.stdout = TextRedirector(self.output_text)
        sys.stderr = TextRedirector(self.output_text)



        def start_create_table():
                # 开一个新线程去执行 create_table
                t = threading.Thread(target=create_table, args=(
                    self.db_type.get(),
                    self.host_entry.get(),
                    self.user_entry.get(),
                    int(self.port_entry.get()),
                    self.pwd_entry.get(),
                    self.dbname_entry.get(),
                    self.schema_entry.get(),
                    self.output_redirector,
                    int(self.rows_entry.get()),
                    self.table_prefix_entry.get()
                ))
                t.start()
            

        ttk.Button(button_frame, text="建造大表环境",
                   command=start_create_table).pack(side=tk.LEFT, padx=5)
    
    def toggle_encrypt_type(self, frame):
            # 先移除已有的加密类型Checkbutton
            for cb in getattr(self, 'encrypt_type_checkbuttons', []):
                cb.grid_forget()
            self.encrypt_type_checkbuttons = []
            if self.build_encrypt_data_var.get():
                md5_cb = ttk.Checkbutton(frame, text="MD5", variable=self.encrypt_md5_var)
                aes_cb = ttk.Checkbutton(frame, text="AES.CDB", variable=self.encrypt_aes_var)
                sha256_cb = ttk.Checkbutton(frame, text="SHA256", variable=self.encrypt_sha256_var)
                md5_cb.grid(row=4, column=4)
                aes_cb.grid(row=4, column=5)
                sha256_cb.grid(row=4, column=6)
                self.encrypt_type_checkbuttons = [md5_cb, aes_cb, sha256_cb]

    def random_date(self):
        start_date = datetime.date(2020, 1, 1)
        end_date = datetime.date(2025, 1, 1)
        delta_days = (end_date - start_date).days
        return start_date + datetime.timedelta(days=random.randint(0, delta_days))

    def random_string(self, length):
        return ''.join(random.choices(string.ascii_letters + string.digits, k=length))

# 构建测试数据类型及格式的方法
    def generate_sql(self):
        global col_type
        try:
            table_count = int(self.table_count_entry.get())
            columns = int(self.columns_entry.get())
            rows = int(self.rows_entry.get())
            table_prefix = self.table_prefix_entry.get()
            col_prefix = self.col_prefix_entry.get()
            db_type = self.db_type.get()
            col_len = self.col_len.get()
            dbname = self.dbname_entry.get()
            schema = self.schema_entry.get()

            output = []
            generated_table_names = set()

            for i in range(1, table_count + 1):
                if self.random_name_var.get():
                    while True:
                        random_suffix = ''.join(random.choices(string.ascii_lowercase, k=6))
                        table_name = table_prefix + random_suffix
                        if table_name not in generated_table_names:
                            generated_table_names.add(table_name)
                            break
                else:
                    table_name = f"{table_prefix}{str(i).zfill(5)}" if self.auto_number_var.get() else table_prefix
                    if table_name in generated_table_names:
                        messagebox.showerror("错误", f"检测到重复的表名: {table_name}")
                        return
                    generated_table_names.add(table_name)

                cols = []
                col_types = []

                for j in range(1, columns + 1):
                    col_name = f"{col_prefix}{j}"
                    if db_type == 'Oracle':
                        col_type = random.choice([f"VARCHAR2({col_len})", "NUMBER", f"char({col_len})"])
                    elif db_type == 'MySQL':
                        col_type = random.choice([f"VARCHAR({col_len})", f"BIGINT(255)"])
                    elif db_type == 'PostgreSQL':
                        col_type = random.choice([f"VARCHAR({col_len})", "BIGSERIAL"])
                    elif db_type == 'GoldenDB':
                        col_type = random.choice([f"VARCHAR({col_len})", "BIGINT(255)"])
                    elif db_type == 'OceanBase':
                        col_type = random.choice([f"VARCHAR({col_len})", "BIGINT(255)"])
                    elif db_type == 'TDSQL':
                        col_type = random.choice([f"VARCHAR({col_len})", "BIGINT(255)"])
                    elif db_type == 'DM':

                        col_type = random.choice([f"VARCHAR({col_len})", "NUMBER"])
                    else:
                        col_type = f"VARCHAR({col_len})"

                    col_types.append(col_type)
                    col_def = f"{col_name} {col_type}"

                    if j == 1 and self.add_pk_var.get():
                        col_def += " PRIMARY KEY"
                    if self.show_type_comment_var.get():
                        col_def += f" -- {col_type}"
                    cols.append(col_def)
                if db_type == 'Oracle':
                    output.append(f"CREATE TABLE {schema}.{table_name} (\n" + ",\n".join(cols) + "\n);")
                elif db_type == 'MySQL':
                    output.append(f"CREATE TABLE {table_name} (\n" + ",\n".join(cols) + "\n);")
                elif db_type == 'PostgreSQL':
                    output.append(f"CREATE TABLE public.{table_name} (\n" + ",\n".join(cols) + f"\n);")
                elif db_type == 'GoldenDB':
                    output.append(f"CREATE TABLE {table_name} (\n" + ",\n".join(cols) + f"\n) distributed by hash({cols[0].split()[0]})(g1,g2);")
                elif db_type == 'OceanBase':
                    output.append(f"CREATE TABLE {table_name} (\n" + ",\n".join(cols) + f"\n) partition by hash({cols[0].split()[0]}) PARTITIONS 3;")
                elif db_type == 'DM':
                    output.append(f"CREATE TABLE {dbname}.{table_name} (\n" + ",\n".join(cols) + f"\n);")

                pk_col_type = col_types[0]
                pk_used_values = set()
                pk_counter = 1

                for r in range(rows):
                    values = []
                    for c in range(columns):
                        dtype = col_types[c].upper()

                        # 主键处理
                        if c == 0 and self.add_pk_var.get():
                            if "CHAR" in pk_col_type or "VARCHAR" in pk_col_type:
                                while True:
                                    pk_val = str(uuid.uuid4())
                                    if pk_val not in pk_used_values:
                                        pk_used_values.add(pk_val)
                                        values.append(f"'{pk_val}'")
                                        break
                            elif "NUMBER" in pk_col_type or "BIGINT" or "BIGSERIAL" in pk_col_type:
                                # values.append(str(pk_counter))
                                # pk_counter += 1
                                values.append(str(random.randint(1, 99999999999)))
                            elif "DATE" in pk_col_type or "TIME" in pk_col_type:
                                values.append(f"'{self.random_date()}'")
                            else:
                                values.append("NULL")
                        else:
                            # 非主键字段处理
                            if "CHAR" in dtype or "VARCHAR" in dtype:
                                if self.encrypt_md5_var.get():
                                    values.append(f"'{hashlib.md5(str(random.randint(1, 99999999999)).encode()).hexdigest()}'")
                                elif self.encrypt_aes_var.get():
                                    data_to_encrypt = str(random.randint(1, 99999999999))
                                    # AES加密，使用16字节密钥和PKCS7填充
                                    key = "1234567890abcdef"  # 16字节密钥，实际可让用户输入
                                    iv = b"abcdef1234567890"  # 16字节IV，实际可随机生成并保存
                                    cipher = AES.new(key.encode(), AES.MODE_CBC, iv)
                                    
                                    padded_data = pad(data_to_encrypt.encode(), AES.block_size)
                                    encrypted_bytes = cipher.encrypt(padded_data).hex()
                                    # messagebox.showinfo("提示", f"AES密钥为{key}, IV为{iv.hex()}，请保存密钥和IV")
                                    values.append(f"'{encrypted_bytes}'")
                                else:
                                    values.append(f"'{self.random_string(8)}'")
                            elif "NUMBER" in dtype or "BIGINT" or "BIGSERIAL" in dtype:
                                values.append(str(random.randint(1, 99999999999)))
                            elif "DATE" in dtype or "TIME" in dtype:
                                if db_type == 'Oracle':
                                    values.append(f"TO_DATE('{self.random_date()}', 'YYYY-MM-DD')")
                                else:
                                    values.append(f"'{self.random_date()}'")
                            else:
                                pass
                    if db_type == 'PostgreSQL':
                        output.append(f"INSERT INTO public.{table_name} VALUES ({', '.join(values)});")
                    elif db_type == 'Oracle':
                        output.append(f"INSERT INTO {schema}.{table_name} VALUES ({', '.join(values)});")
                    else:
                        output.append(f"INSERT INTO {dbname}.{table_name} VALUES ({', '.join(values)});")
                output.append("\n")

            self.output_text.delete(1.0, tk.END)
            self.output_text.insert(tk.END, '\n'.join(output))

        except ValueError:
            messagebox.showerror("错误", "请输入有效的数字")

    def export_sql(self):
        sql = self.output_text.get(1.0, tk.END).strip()
        if not sql:
            messagebox.showinfo("提示", "没有生成任何SQL")
            return
        filepath = filedialog.asksaveasfilename(defaultextension=".sql", filetypes=[("SQL 文件", "*.sql")])
        if filepath:
            with open(filepath, 'w', encoding='utf-8') as f:
                f.write(sql)
            messagebox.showinfo("成功", f"SQL 已导出到\n{filepath}")
#
    def clear_output(self):
        self.output_text.delete(1.0, tk.END)

    def simulate_execute(self):
        sql = self.output_text.get(1.0, tk.END).strip()
        if not sql:
            messagebox.showinfo("提示", "没有SQL可以执行")
            return

        db_type = self.db_type.get()
        host = self.host_entry.get()
        port = int(self.port_entry.get())
        dbname = self.dbname_entry.get()
        user = self.user_entry.get()
        password = self.pwd_entry.get()

        try:

            if db_type == "Oracle":
                dsn = cx_Oracle.makedsn(host, port, dbname)
                conn = cx_Oracle.connect(user=user, password=password, dsn=dsn)
            elif db_type == "MySQL":
                conn = pymysql.connect(host=host, port=port, user=user, password=password, db=dbname)
            elif db_type == "PostgreSQL":
                conn = psycopg2.connect(host=host, port=port, user=user, password=password, dbname=dbname)
            elif db_type == "GoldenDB":
                conn = pymysql.connect(host=host, port=port, user=user, password=password, db=dbname)
            elif db_type == "TDSQL":
                conn = pymysql.connect(host=host, port=port, user=user, password=password, db=dbname)
            elif db_type == "OceanBase":
                conn = pymysql.connect(host=host, port=port, user=user, password=password, db=dbname)
            elif db_type == "DM":
                import jaydebeapi
                # 设置数据库连接信息
                jdbc_url = f"jdbc:dm://{host}:{port}"  # 5236为实例端口号，每个实例对应的端口号不同，这里通过指定端口号来连接到不同实例中，不需要实例名
                jdbc_driver = "dm.jdbc.driver.DmDriver"
                jdbc_jar_path = r"./jdbc_driver/DmJdbcDriver18.jar"  # 替换成实际的 JDBC 驱动 JAR 文件路径
                conn = jaydebeapi.connect(jdbc_driver, jdbc_url, [user, password], jdbc_jar_path)
            else:
                raise ValueError("不支持的数据库类型")

            cursor = conn.cursor()
            for sql_line in sql.split(';'):
                if sql_line.strip():
                    # cursor.executem(sql_line)
                    cursor.execute(sql_line)

            conn.commit()
            messagebox.showinfo("成功", "SQL执行成功")

        except Exception as e:
            messagebox.showerror("错误", f"执行失败: {str(e)}")

    def test_connection(self):
        host = self.host_entry.get()
        port = int(self.port_entry.get())
        dbname = self.dbname_entry.get()
        user = self.user_entry.get()
        password = self.pwd_entry.get()

        db_type = self.db_type.get()
        try:
            if db_type == "Oracle":
                conn = cx_Oracle.connect(
                    user=user,
                    password=password,
                    dsn=f"{host}:{port}/{dbname}"
                )
            elif db_type == "MySQL":
                conn = pymysql.connect(host=host, port=port, user=user, password=password, db=dbname)
            elif db_type == "GoldenDB":
                conn = pymysql.connect(host=host, port=port, user=user, password=password, db=dbname)
            elif db_type == "TDSQL":
                conn = pymysql.connect(host=host, port=port, user=user, password=password, db=dbname)
            elif db_type == "PostgreSQL":
                conn = psycopg2.connect(host=host, port=port, user=user, password=password, dbname=dbname)
            elif db_type == "OceanBase":
                conn = pymysql.connect(host=host, port=port, user=user, password=password, db=dbname)
            elif db_type == "DM":
                import jaydebeapi
                # 设置数据库连接信息
                jdbc_url = f"jdbc:dm://{host}:{port}"  # 5236为实例端口号，每个实例对应的端口号不同，这里通过指定端口号来连接到不同实例中，不需要实例名
                jdbc_driver = "dm.jdbc.driver.DmDriver"
                jdbc_jar_path = r"./jdbc_driver/DmJdbcDriver18.jar"  # 替换成实际的 JDBC 驱动 JAR 文件路径
                conn = jaydebeapi.connect(jdbc_driver, jdbc_url, [user, password], jdbc_jar_path)
            else:
                raise ValueError("不支持的数据库类型")

            messagebox.showinfo("连接成功", "数据库连接成功！")
        except Exception as e:
            messagebox.showerror("连接失败", f"无法连接到数据库: {str(e)}")


if __name__ == "__main__":
    root = tk.Tk()
    app = TableGeneratorApp(root)
    root.mainloop()

    def toggle_encrypt_type(self, frame):
        # 先移除已有的加密类型Checkbutton
        for cb in getattr(self, 'encrypt_type_checkbuttons', []):
            cb.grid_forget()
        self.encrypt_type_checkbuttons = []
        if self.build_encrypt_data_var.get():
            md5_cb = ttk.Checkbutton(frame, text="MD5", variable=self.encrypt_md5_var)
            aes_cb = ttk.Checkbutton(frame, text="AES.CDB", variable=self.encrypt_aes_var)
            sha256_cb = ttk.Checkbutton(frame, text="SHA256", variable=self.encrypt_sha256_var)
            md5_cb.grid(row=4, column=4)
            aes_cb.grid(row=4, column=5)
            sha256_cb.grid(row=4, column=6)
            self.encrypt_type_checkbuttons = [md5_cb, aes_cb, sha256_cb]